import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 1. 数据预处理和加载
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL图像转换为Tensor，并自动归一化到[0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
])

# 下载并加载训练集和测试集
train_dataset = datasets.MNIST(
    root='./data', 
    train=True, 
    download=True, 
    transform=transform
)

test_dataset = datasets.MNIST(
    root='./data', 
    train=False, 
    download=True, 
    transform=transform
)

# 创建数据加载器
train_loader = DataLoader(
    train_dataset, 
    batch_size=64, 
    shuffle=True  # 训练时打乱数据
)

test_loader = DataLoader(
    test_dataset, 
    batch_size=1000, 
    shuffle=False  # 测试时不需要打乱
)

# 2. 定义神经网络模型
class MNISTClassifier(nn.Module):
    def __init__(self):
        super(MNISTClassifier, self).__init__()
        # 第一个卷积层：输入通道1(灰度图)，输出通道32，卷积核3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        # 第二个卷积层：输入通道32，输出通道64，卷积核3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 最大池化层：2x2窗口
        self.pool = nn.MaxPool2d(2, 2)
        # Dropout层，防止过拟合
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # 经过两次池化后尺寸为7x7
        self.fc2 = nn.Linear(128, 10)  # 输出10个类别（数字0-9）
    
    def forward(self, x):
        # 第一个卷积块：卷积 -> ReLU -> 池化 -> Dropout
        x = self.pool(F.relu(self.conv1(x)))
        x = self.dropout1(x)
        
        # 第二个卷积块：卷积 -> ReLU -> 池化 -> Dropout
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout1(x)
        
        # 展平特征图
        x = x.view(-1, 64 * 7 * 7)
        
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        
        return F.log_softmax(x, dim=1)  # 输出对数概率

# 3. 初始化模型、损失函数和优化器
model = MNISTClassifier().to(device)
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

# 4. 训练函数
def train(model, device, train_loader, optimizer, epoch):
    model.train()  # 设置为训练模式
    train_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 反向传播
        loss.backward()
        optimizer.step()
        
        # 统计信息
        train_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        # 每100个batch打印一次进度
        if batch_idx % 100 == 0:
            print(f'训练周期: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\t损失: {loss.item():.6f}')
    
    # 计算训练准确率
    accuracy = 100. * correct / total
    avg_loss = train_loss / len(train_loader)
    return avg_loss, accuracy

# 5. 测试函数
def test(model, device, test_loader):
    model.eval()  # 设置为评估模式
    test_loss = 0
    correct = 0
    
    with torch.no_grad():  # 测试时不需要计算梯度
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)  # 获取预测结果
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    test_loss /= len(test_loader)
    accuracy = 100. * correct / len(test_loader.dataset)
    
    print(f'\n测试集: 平均损失: {test_loss:.4f}, 准确率: {correct}/{len(test_loader.dataset)} '
          f'({accuracy:.2f}%)\n')
    
    return test_loss, accuracy

# 6. 可视化一些样本
def visualize_samples(loader, num_samples=6):
    dataiter = iter(loader)
    images, labels = next(dataiter)
    
    fig, axes = plt.subplots(1, num_samples, figsize=(12, 2))
    for i in range(num_samples):
        axes[i].imshow(images[i][0], cmap='gray')
        axes[i].set_title(f'标签: {labels[i].item()}')
        axes[i].axis('off')
    plt.tight_layout()
    plt.show()

# 可视化训练样本
print("训练样本示例:")
visualize_samples(train_loader)

# 7. 训练循环
train_losses = []
train_accuracies = []
test_losses = []
test_accuracies = []

epochs = 5
for epoch in range(1, epochs + 1):
    train_loss, train_acc = train(model, device, train_loader, optimizer, epoch)
    test_loss, test_acc = test(model, device, test_loader)
    
    # 保存训练过程数据
    train_losses.append(train_loss)
    train_accuracies.append(train_acc)
    test_losses.append(test_loss)
    test_accuracies.append(test_acc)

# 8. 绘制训练曲线
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), train_losses, 'b-', label='训练损失')
plt.plot(range(1, epochs+1), test_losses, 'r-', label='测试损失')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('训练和测试损失')

plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), train_accuracies, 'b-', label='训练准确率')
plt.plot(range(1, epochs+1), test_accuracies, 'r-', label='测试准确率')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()
plt.title('训练和测试准确率')

plt.tight_layout()
plt.show()

# 9. 在单个样本上测试
def predict_single_sample(model, device, sample_data, sample_label):
    model.eval()
    with torch.no_grad():
        output = model(sample_data.unsqueeze(0).to(device))
        pred = output.argmax(dim=1, keepdim=True)
        prob = F.softmax(output, dim=1)
        
    plt.imshow(sample_data[0], cmap='gray')
    plt.title(f'真实标签: {sample_label}, 预测: {pred.item()}, 置信度: {prob[0][pred].item():.4f}')
    plt.axis('off')
    plt.show()

# 测试单个样本
sample_data, sample_label = test_dataset[0]
predict_single_sample(model, device, sample_data, sample_label)

print("训练完成！")